package com.steven.logserver.utils;

import org.elasticsearch.action.admin.indices.create.CreateIndexRequestBuilder;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequestBuilder;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryAction;
import org.elasticsearch.index.reindex.DeleteByQueryRequestBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.BucketOrder;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.avg.Avg;
import org.elasticsearch.search.aggregations.metrics.max.Max;
import org.elasticsearch.search.aggregations.metrics.min.Min;
import org.elasticsearch.search.aggregations.metrics.sum.Sum;
import org.elasticsearch.search.sort.SortOrder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import javax.annotation.PostConstruct;
import java.util.*;

/**
 * ES 工具类（注入组件 + 静态方法 组合方式的工具类）
 *
 * @author qinxiaohu
 */
@Component
public class ElasticsearchUtil {

    private static final Logger LOGGER = LoggerFactory.getLogger(ElasticsearchUtil.class);

    @Autowired
    private TransportClient transportClient;
    private static TransportClient client;

    /**
     * id 关键字
     */
    private static final String ID = "id";

    /**
     * 聚合查询一些关键字：
     *  termsName（聚合字段。不可缺省）、size（返回数量。可缺省）、order（是否按集合字段排序。可缺省）
     *  、subAgg（其他字段聚合函数，存放聚合字段和聚合函数的map。可缺省）
     */
    private static final String AGG_TERMS_NAME = "termsName";
    private static final String AGG_ORDER = "order";
    private static final String AGG_SUB_AGG = "subAgg";
    private static final String AGG_RETURN_SIZE = "size";

    /**
     * 聚合函数的求值方式关键字
     */
    private static final String SUM = "sum";
    private static final String AVG = "avg";
    private static final String MAX = "max";
    private static final String MIN = "min";

    @PostConstruct
    public void init() {

        client = this.transportClient;
    }

    //------------ 索引和类型相关操作 ---------------

    /**
     * 判断索引是否存在
     *
     * @param index     索引
     * @return boolean
     */
    public static boolean isIndexExist(String index) {

        return client.admin().indices().prepareExists(index).execute().actionGet().isExists();
    }

    /**
     * 删除索引
     *
     * @param index     索引
     * @return boolean
     */
    public static boolean deleteIndex(String index) {

        return isIndexExist(index) && client.admin().indices().prepareDelete(index).execute().actionGet().isAcknowledged();
    }

    /**
     * 新增索引
     *
     * @param index     索引
     * @return boolean
     */
    public static boolean addIndex(String index) {

        return !isIndexExist(index) && client.admin().indices().prepareCreate(index).execute().actionGet().isAcknowledged();
    }

    /**
     * 判断index下指定type是否存在
     *
     * @param index     索引
     * @param type      类型
     * @return boolean
     */
    public static boolean isTypeExist(String index, String type) {

        return isIndexExist(index) &&
                client.admin().indices().prepareTypesExists(index).setTypes(type).execute().actionGet().isExists();
    }

    /**
     * 新增类型和mapping
     *
     * @param index             索引
     * @param type              类型
     * @param mapping           映射
     * @param setting           索引的setting （包含分片 ，备份之类的）
     *
     * @return
     */
    public static boolean addMapping(String index, String type, XContentBuilder mapping, Settings.Builder setting) {

        // 创建索引映射,相当于创建数据库中的表操作
        CreateIndexRequestBuilder cib = client.admin().indices().prepareCreate(index);
        // 设置mapping
        cib.addMapping(type, mapping);
        // 创建索引的setting （包含分片，备份之类的）
        cib.setSettings(setting);

        return cib.execute().actionGet().isAcknowledged();
    }

    //------------ 常规的增删查改操作 ---------------

    /**
     * 数据添加，正定ID
     *
     * @param jsonObject 					要增加的数据
     * @param index      					索引，类似数据库
     * @param type       					类型，类似表
     * @param id         					数据ID
     * @return String
     */
    public static String addData(Map<String,Object> jsonObject, String index, String type, String id) {

        // 如果id存在，则直接使用
        if(jsonObject.get(ID) != null){
            id = jsonObject.get(ID).toString();
        }
        IndexResponse response = client.prepareIndex(index, type, id).setSource(jsonObject).get();
        LOGGER.info("addData result is status:{},id:{}", response.status().getStatus(), response.getId());

        return response.getId();
    }

    /**
     * 批量添加，id 随机（可以根据业务需要修改策略）
     *
     * @param jsonObject 							数据集合
     * @param index									索引
     * @param type									类型
     * @return										返回导入数据数量
     * @throws InterruptedException
     */
    public static Integer bulkAddData(List<Map<String,Object>> jsonObject, String index, String type) throws InterruptedException{

        BulkRequestBuilder bulkRequest = client.prepareBulk();

        Integer sum = 0;
        for(Map<String,Object > object : jsonObject){
            String id ;
            if(object.get(ID) != null){
                id = object.get(ID).toString();
            }else{
                id = UUID.randomUUID().toString().replaceAll("-", "").toUpperCase();
            }
            bulkRequest.add(client.prepareIndex(index,type,id).setSource(object));
            sum++;
            // 50000条提交一次
            if(sum > 0 && sum % 50000 == 0){
                bulkRequest.execute().actionGet();
                bulkRequest = client.prepareBulk();
                // 线程休息5秒,给出时间让ES服务端处理
                Thread.sleep(5000);
            }
        }
        bulkRequest.execute().actionGet();
        LOGGER.info("bulkAddData result is index:{},type:{},sumNum:{}",index, type,sum);

        return sum;
    }

    /**
     * 数据添加
     *
     * @param jsonObject 						要增加的数据
     * @param index      						索引，类似数据库
     * @param type       						类型，类似表
     * @return String
     */
    public static String addData(Map<String,Object> jsonObject, String index, String type) {

        String id = UUID.randomUUID().toString().replaceAll("-", "").toUpperCase();
        return addData(jsonObject, index, type, id);
    }

    /**
     * 通过ID删除数据
     *
     * @param index 							索引，类似数据库
     * @param type  							类型，类似表
     * @param id    							数据ID
     */
    public static void deleteDataById(String index, String type, String id) {

        DeleteResponse response = client.prepareDelete(index, type, id).execute().actionGet();

        LOGGER.info("deleteDataById result is status:{},id:{}",response.status().getStatus(), response.getId());
    }

    /**
     * 通过条件删除数据
     *
     * @param index								索引，类似数据库
     * @param builder							条件
     * @return Long								返回删除总条数
     */
    public static Long deleteDataByParam(String index, BoolQueryBuilder builder) {

        DeleteByQueryRequestBuilder requestBuilder =
                DeleteByQueryAction.INSTANCE.newRequestBuilder(client).filter(builder).source(index);

        BulkByScrollResponse response = requestBuilder.get();

        return response.getDeleted();
    }

    /**
     * 批量删除数据
     *
     * @param index								索引
     * @param type								类型
     * @param idList							需要删除数据的id列表
     * @throws InterruptedException
     */
    public static void bulkDeleteData(String index, String type, List<String> idList) throws InterruptedException {

        BulkRequestBuilder bulkRequest = client.prepareBulk();

        for(int i=0; i<idList.size(); i++){

            bulkRequest.add(client.prepareDelete(index, type, idList.get(i)));
            if( i > 0 && i % 50000 == 0 ){
                bulkRequest.execute().actionGet();
                bulkRequest = client.prepareBulk();
                //  线程休息5秒,给出时间让ES服务端处理
                Thread.sleep(5000);
            }
        }
        bulkRequest.execute().actionGet();

        LOGGER.info("bulkDeleteData result is sumNum:{}",idList.size());
    }

    /**
     * 通过ID 更新数据
     *
     * TIPS：通过新增也可以实现新增，如果指定数据的id在中存在，那么就会覆盖id对应的原来数据
     *
     * @param jsonObject 						要增加的数据
     * @param index      						索引，类似数据库
     * @param type       						类型，类似表
     * @param id        	 					数据ID
     * @return
     */
    public static Boolean updateDataById(Map<String,Object> jsonObject, String index, String type, String id) {

        try{
            UpdateRequest updateRequest = new UpdateRequest();
            updateRequest.index(index).type(type).id(id).doc(jsonObject);
            client.update(updateRequest);
            return Boolean.TRUE;
        }catch (Exception e){
            LOGGER.error("updateDataById error:",e);
            return Boolean.FALSE;
        }
    }

    /**
     * 通过ID获取数据
     *
     * @param index  							索引，类似数据库
     * @param type   							类型，类似表
     * @param id     							数据ID
     * @param fields 							需要显示的字段，逗号分隔（缺省为全部字段）
     * @return Map<String, Object>
     */
    public static Map<String, Object> searchDataById(String index, String type, String id, String fields) {

        GetRequestBuilder getRequestBuilder = client.prepareGet(index, type, id);

        if (!StringUtils.isEmpty(fields)) {
            getRequestBuilder.setFetchSource(fields.split(","), null);
        }
        GetResponse getResponse =  getRequestBuilder.execute().actionGet();

        return getResponse.getSource();
    }

    //------------ 通过DSL查询操作 ---------------

    /**
     * DSL查询结果（不做分页处理，不排序）
     *
     * @param index								索引（类似数据库）
     * @param type								类型（类似表）
     * @param builder							es查询查询条件需要的QueryBuilder
     * @param aggMap							需要聚合的数据
     * @return List<Map<String, Object>>
     */
    public static List<Map<String, Object>> aggDataByQueryNotPage(String index, String type
            , BoolQueryBuilder builder, Map<String,Object> aggMap){

        return searchAggDataByQuery( index, type, builder,null,null,null,null, aggMap);
    }

    /**
     * DSL查询结果（不做分页处理，不排序，但是指定需要返回的指定字段）
     *
     * @param index								索引（类似数据库）
     * @param type								类型（类似表）
     * @param builder							es查询需要的QueryBuilder
     * @param aggMap							聚合数据
     * @param includes							需要返回的指定字段
     * @return List<Map<String, Object>>
     */
    public static List<Map<String, Object>> aggDataByQueryNotPageIncludes(String index, String type, String[] includes
            , BoolQueryBuilder builder, Map<String, Object> aggMap){

        return searchAggDataByQuery( index, type, builder,null,null,includes,null, aggMap);
    }

    /**
     * DSL查询结果（不做分页处理，但排序）
     *
     * @param index								索引（类似数据库）
     * @param type								类型（类似表）
     * @param builder							es查询需要的QueryBuilder
     * @param aggMap							聚合数据
     * @param sort								排序规则
     * @return List<Map<String, Object>>
     */
    public static List<Map<String, Object>> aggDataByQuerySort(String index, String type, LinkedHashMap<String
            , SortOrder> sort, BoolQueryBuilder builder,Map<String,Object> aggMap){

        return searchAggDataByQuery( index,type,builder,sort,null,null, null, aggMap);
    }

    /**
     *
     * DSL查询结果（带分页）
     *
     * TIPS：使用from + size分页方式存在性能问题，需要使用快照来解决
     *
     * @param index         索引
     * @param type          类型（多个用“，”隔开）
     * @param builder       查询条件
     * @param sort          排序
     * @param from          下标
     * @param aggMap        聚合条件
     * @param size          获取数据数量
     * @param includes      指定返回字段（为null返回所有）
     *
     * @return List<Map<String, Object>>
     */
    @SuppressWarnings("unchecked")
    public static List<Map<String, Object>> searchAggDataByQuery(String index, String type, BoolQueryBuilder builder
            , LinkedHashMap<String, SortOrder> sort, Integer from, String[] includes, Integer size, Map<String,Object> aggMap) {

        SearchRequestBuilder search = client.prepareSearch(index);
        // 1. 设置类型
        if (!StringUtils.isEmpty(type)) {
            search.setTypes(type.split(","));
        }
        // 2. 添加排序
        if(sort != null){
            sort.forEach(search::addSort);
        }
        // 3. 设置下标和步长。TIPS: 一般搜索分页基于form+size实现。size超过1万数据就会抛出异常，最好限制在1万之内。当所请求的数据总量大于1w时可用scroll来代替from+size（待做）
        if(from != null && from > 0) {
            search.setFrom((from - 1) * size);
        }
        if(size != null && size > 0) {
            search.setSize(size);
        }
        // 4. 指定返回的字段
        // ① setFetchSource(boolean fetch)  用于指定是否返回字段，默认为true。如果为false，则getSource返回null；
        // ② setFetchSource(String include, String exclude) 用于指定包含的一个字段和排除的一个字段，实际会返回include字段及其对应的值；
        // ③ setFetchSource(String[] includes, String[] excludes) 用于指定包含的多个字段和排除的多个字段；
        if(includes != null && includes.length > 0){
            search.setFetchSource(includes, null);
        }
        // 5. 根据aggMap是否为空，来决定为聚合查询还是DSL查询操作
        SearchResponse response;
        if(aggMap == null){
            response= search.setQuery(builder).get();
        }else{
            // 5.1 聚合查询没必要返回查询字段，所以设置不返回查询字段；因为没必要获取非聚合数据，所以设置聚合查询的hits返回结果为0
            search.setFetchSource(false);
            search.setSize(0);
            String termsName = aggMap.get(AGG_TERMS_NAME).toString();
            if(org.springframework.util.StringUtils.isEmpty(termsName)){
                throw new RuntimeException("聚合字段缺失");
            }
            TermsAggregationBuilder aggregationBuilder = AggregationBuilders.terms(termsName).field(termsName);
            // 5.2 默认聚合后返回条数为10，可以指定返回数据数量
            Object size1 = aggMap.get(AGG_RETURN_SIZE);
            if(size1 != null && Integer.valueOf(size1.toString()) > 0 ){
                aggregationBuilder.size(Integer.valueOf(size1.toString()));
            }
            // 5.3 默认不排序，可以指定聚合字段termsName的排序方式
            Object order = aggMap.get(AGG_ORDER);
            if(order != null){
                BucketOrder.aggregation(termsName, (Boolean)order);
            }
            // 5.4 设置其他字段需要的聚合函数，例如avg()、sum()、max()、min()等
            Object subAgg = aggMap.get(AGG_SUB_AGG);
            if(subAgg != null){
                Map<String, String> map = (Map<String, String>) subAgg;
                map.forEach((key, value) -> {
                    switch (value){
                        case SUM:
                            aggregationBuilder.subAggregation(AggregationBuilders.sum(SUM + "_" + key).field(key));
                            break;
                        case AVG:
                            aggregationBuilder.subAggregation(AggregationBuilders.avg(AVG + "_"  + key).field(key));
                            break;
                        case MAX:
                            aggregationBuilder.subAggregation(AggregationBuilders.max(MAX + "_"  + key).field(key));
                            break;
                        case MIN:
                            aggregationBuilder.subAggregation(AggregationBuilders.min(MIN + "_"  + key).field(key));
                            break;
                        default:
                            throw new RuntimeException("不存在的聚合函数类型");
                    }
                });
            }
            response = search.setQuery(builder).addAggregation(aggregationBuilder).get();
            LOGGER.info("response status is {}",response.status());
        }
        if(response == null){
            return Collections.emptyList();
        }
        // 6. 普通DSL查询直接返回值，聚合查询需要获取聚合数据
        List<Map<String, Object>> sourceList = new ArrayList<>();
        if(aggMap == null){
            SearchHit[] hits = response.getHits().getHits();
            if(hits == null || hits.length == 0){
                return Collections.emptyList();
            }
            for (SearchHit searchHit : hits) {
                sourceList.add(searchHit.getSourceAsMap());
            }
        }else{
            Terms termsName = response.getAggregations().get(aggMap.get(AGG_TERMS_NAME).toString());
            if(termsName == null){
                return Collections.emptyList();
            }
            for (Terms.Bucket bucket : termsName.getBuckets()) {
                Map<String, Object> returnMap = new HashMap<>(16);
                // 6.1 得到聚合字段的关键字
                Object o =bucket.getKey();
                // 6.2 得到聚合字段值的count值
                long count = bucket.getDocCount();
                returnMap.put("key", o);
                returnMap.put("count", count);
                // 6.3 得到聚合函数Map，遍历得到集合后的值
                Map<String, String> subAggMap = (Map<String, String>) aggMap.get(AGG_SUB_AGG);
                subAggMap.forEach((key,value) -> {
                    switch (value){
                        case SUM:
                            Sum sum = bucket.getAggregations().get(SUM + "_" + key);
                            returnMap.put(SUM + "_" + key, sum.getValue());
                            break;
                        case AVG:
                            Avg avg = bucket.getAggregations().get(AVG + "_" + key);
                            returnMap.put(AVG + "_"+ key, avg.getValue());
                            break;
                        case MAX:
                            Max max = bucket.getAggregations().get(MAX + "_" + key);
                            returnMap.put(MAX + "_" + key, max.getValue());
                            break;
                        case MIN:
                            Min min = bucket.getAggregations().get(MIN + "_" + key);
                            returnMap.put(MIN + "_" + key, min.getValue());
                            break;
                        default:
                            throw new RuntimeException("不存在的聚合函数类型");
                    }
                });
                sourceList.add(returnMap);
            }
        }
        return sourceList;
    }

}
